import numpy as np
import torch
from torch import nn
from torch.autograd import Function

class Conv2d(Function):
    @staticmethod
    def forward(ctx,input,weight,bias,stride=1,padding=0):
        N,C_in,H,W=input.shape
        C_out,_,kH,KW=weight.shape
        H_out=



model=SimpleCNN()
criterion=nn.CrossEntropyLoss()
optimizer=torch.optim.SGD(model.parameters(),lr=0.01)
x=torch.rand(16,3,640*640)
y=torch.randint(0,2,(16,))
for epoch in range(10):
    out=model(x)
    loss=criterion(out,y)
    